from torch.utils.data import Dataset
from torch import cat as T_cat
from csv import reader as CSV_reader
from os.path import join as o_p_join
from torchvision.transforms import Compose, ToTensor, Normalize
from PIL import Image
from random import randrange


PAIR_DIR = 'Protocol'
PAIR_FILE = 'Bmp_list.txt'
LIST_FILE = 'Data/list_name.txt'
NUM_CLASSES = 500

VCFP_DIR = 'Protocol/Split/FF'
VCFP_BATCH = 10
VCFP_BSIZE = 350


class CFP_Dset:
    def __init__(self, root_dir):
        # 创建类型列表
        self.classes = []
        with open(o_p_join(root_dir, LIST_FILE)) as csv_f:
            csv_r = CSV_reader(csv_f)
            for _ in range(NUM_CLASSES):
                row_data = next(csv_r)
                self.classes.append(row_data[0])
        # 统计各项数据列表
        self.root_dir = o_p_join(root_dir, PAIR_DIR)
        self.item_dir = []
        self.item_count = 0
        with open(o_p_join(self.root_dir, PAIR_FILE)) as csv_f:
            csv_r = CSV_reader(csv_f, delimiter=' ')
            for _ in range(NUM_CLASSES):
                for _ in range(10):
                    row_data = next(csv_r)
                    self.item_count += 1
                    self.item_dir.append(row_data[1])
        # 确定数据预处理方式
        self.item_trans = Compose([ToTensor(), Normalize([0.5, 0.5, 0.5], [0.2, 0.2, 0.2])])

    def __getitem__(self, idx):
        img = Image.open(o_p_join(self.root_dir, self.item_dir[idx]))
        return self.item_trans(img)

    def __len__(self):
        return self.item_count


class CFP_Random(Dataset):
    def __init__(self, root_dir):
        super(CFP_Random, self).__init__()
        self.dset_src = CFP_Dset(root_dir)
        self.a_bias = 0
        self.p_bias = 0
        self.n_bias = 0
        self.p_class = 0
        self.n_class = 0
        self.a_idx = 0
        self.p_idx = 0
        self.n_idx = 0
        self.p_str = ''
        self.n_str = ''
    
    def __getitem__(self, _):
        # 获取人脸所属类型
        self.p_class = randrange(0, NUM_CLASSES)
        self.n_class = randrange(0, NUM_CLASSES - 1)
        if self.n_class == self.p_class:
            self.n_class += 1
        self.p_str = self.dset_src.classes[self.p_class]
        self.n_str = self.dset_src.classes[self.n_class]
        # 获取人脸下标
        self.a_bias = randrange(0, 10)
        self.p_bias = randrange(0, 9)
        if self.p_bias == self.a_bias:
            self.p_bias += 1
        self.n_bias = randrange(0, 10)
        self.a_idx = self.p_class * 10 + self.a_bias
        self.p_idx = self.p_class * 10 + self.p_bias
        self.n_idx = self.n_class * 10 + self.n_bias
        return T_cat((self.dset_src[self.a_idx].unsqueeze(0),
                        self.dset_src[self.p_idx].unsqueeze(0),
                        self.dset_src[self.n_idx].unsqueeze(0)), dim=0)

    def __len__(self):
        return len(self.dset_src)


class CFP_Valid(Dataset):
    def __init__(self, root_dir):
        super(CFP_Valid, self).__init__()
        self.dset_src = CFP_Dset(root_dir)
        self.root_dir = o_p_join(root_dir, VCFP_DIR)
        # 数据编号获取
        self.diff_idx_l = []
        self.diff_idx_r = []
        self.same_idx_l = []
        self.same_idx_r = []
        self.item_count = 0
        for i in range(VCFP_BATCH):
            str_i = str(i + 1).zfill(2)
            diff_path = o_p_join(self.root_dir, str_i, 'diff.txt')
            with open(diff_path) as csv_f:
                csv_r = CSV_reader(csv_f)
                for _ in range(VCFP_BSIZE):
                    row_data = next(csv_r)
                    self.diff_idx_l.append(int(row_data[0]) - 1)
                    self.diff_idx_r.append(int(row_data[1]) - 1)
                    self.item_count += 1
            same_path = o_p_join(self.root_dir, str_i, 'same.txt')
            with open(same_path) as csv_f:
                csv_r = CSV_reader(csv_f)
                for _ in range(VCFP_BSIZE):
                    row_data = next(csv_r)
                    self.same_idx_l.append(int(row_data[0]) - 1)
                    self.same_idx_r.append(int(row_data[1]) - 1)
                    self.item_count += 1

    def __getitem__(self, idx):
        block = idx // (2 * VCFP_BSIZE)
        b_off = idx % (2 * VCFP_BSIZE)
        from_same = b_off // VCFP_BSIZE
        t_off = b_off % VCFP_BSIZE
        t_idx = block * VCFP_BSIZE + t_off
        if from_same:
            img_l = self.dset_src[self.same_idx_l[t_idx]]
            img_r = self.dset_src[self.same_idx_r[t_idx]]
        else:
            img_l = self.dset_src[self.diff_idx_l[t_idx]]
            img_r = self.dset_src[self.diff_idx_r[t_idx]]
        return T_cat((img_l.unsqueeze(0), img_r.unsqueeze(0)), dim=0), from_same, block

    def __len__(self):
        return self.item_count
